# Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
import random
from typing import Union, List

from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig


def random_factor_pair(M):
    # 找出所有正因子
    factors = [i for i in range(1, M+1) if M % i == 0]
    # 随机选择一个因子A
    A = random.choice(factors)
    B = M // A
    return A, B

@GENERATOR_REGISTRY.register("generator_ascend_generate_moe_finalize_routing_v2_grad")
class AscendMoeFinalizeRoutingV2(CaseGenerator):
    def __init__(self, config):
        super().__init__(config)

    def after_input_config(
            self,
            index: int,
            input_case: Union[InputCaseConfig, List[InputCaseConfig]]
    ) -> Union[InputCaseConfig, List[InputCaseConfig]]:

        return input_case

    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
        gradY = case_config.inputs[0]
        expandedRowIdx = case_config.inputs[1]
        expandedX = case_config.inputs[2]
        scalesOptional = case_config.inputs[3]
        expertIdxOptional = case_config.inputs[4]
        biasOptional = case_config.inputs[5]
        dropPadMode = case_config.inputs[6]
        activeNum = case_config.inputs[7]
        expertNum = case_config.inputs[8]
        expertCapacity = case_config.inputs[9]
        # dropPadMode
        R, H = gradY.shape[0] , gradY.shape[1]
        E, C = expandedX.shape[0], expandedX.shape[1]
        K = scalesOptional.shape[1]
        if K > E:
            K = E
        scalesOptional.shape[1] = K # K: 为从总的专家E中选出K个专家;  E: expert num，即专家数，E需要大于等于K; 
        expandedRowIdx.shape = [R * K] # 要求是一个1D的Tensor,其shape支持（R * K）
        expandedRowIdx.range_values = [-1, E * C - 1] # dropPadMode参数值为1、3时，Tensor中的值取值范围是[-1, E\*C - 1]。
        expandedX.shape = [E, C, H] # 限制：其shape支持（E，C，H）
        biasOptional.shape = [E, H] # 限制：其shape支持（E，H）
        biasOptional.dtype = expandedX.dtype # 数据类型要求与expandedX一致
        scalesOptional.shape = [R, K] # 限制：其shape支持（NUM\_ROWS，K）
        scalesOptional.dtype = expandedX.dtype # 数据类型要求与expandedX一致
        expertIdxOptional.shape = [R, K] # 限制：其shape支持（NUM\_ROWS，K）
        expertIdxOptional.range_values = [0, E -1] # 限制：Tensor中的值取值范围是[0, E-1]
        expertCapacity.range_values = [C]
        expertNum.range_values = [E]
        return case_config